from functools import cached_property
import io
import logging
import math
from pathlib import Path
import re
from enum import Enum
import zipfile
import polars as pl
import matplotlib.pyplot as plt
from pydantic import BaseModel, Field
from pysam import VariantFile  # pylint: disable=E0611
from .utils import is_active_chrom, is_male, is_sex_chrom

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("ClinCNV")


class Color(Enum):
    BLUE = "#7EA2ED"
    ORANGE = "#FAAA89"
    RED = "#DF8D8F"
    GRAY = "#808080"
    GRAY2 = "#E6E6E6"
    GRAY3 = "#D0D0D0"


Stain2Color = dict(
    zip(
        ["acen", "gneg", "gpos", "gvar", "stalk", "gpos25", "gpos50", "gpos75", "gpos100", "gpos33", "gpos66"],
        ["#808080", "#E6E6E6", "#333333", "#999999", "#BFBFBF", "#B3B3B3", "#808080", "#4D4D4D", "#1A1A1A", "#999999", "#666666"],
    )
)


class Image(BaseModel):
    sample: str
    gender: str
    cytoband_file: Path
    call_vcf_file: Path
    multi_cn_file: Path
    mask_cn_file: Path
    genepred_file: Path
    image_file: Path
    dpi: int = Field(default=300)
    figsize: tuple[int, int] = Field(default=(15, 4))

    @cached_property
    def multi_cn_records(self) -> pl.DataFrame:
        return pl.read_csv(self.multi_cn_file, separator="\t")

    @cached_property
    def mask_cn_records(self) -> pl.DataFrame:
        return pl.read_csv(self.mask_cn_file, separator="\t", null_values="NaN").filter(pl.col("cn").is_not_null())

    @cached_property
    def gene_records(self) -> pl.DataFrame:
        records = pl.read_csv(
            self.genepred_file,
            separator="\t",
            has_header=False,
            columns=[1, 2, 3, 4, 5, 9, 10, 12],
            new_columns=["transcript", "chromosome", "strand", "tx_start", "tx_end", "start", "end", "gene"],
        )
        return (
            records.with_columns(pl.col("start").str.strip_chars(",").str.split(",").cast(pl.List(pl.Int64)).alias("start"))
            .with_columns(pl.col("end").str.strip_chars(",").str.split(",").cast(pl.List(pl.Int64)).alias("end"))
            .select(["chromosome", "gene", "strand", "tx_start", "tx_end", "start", "end"])
        )

    @cached_property
    def cytoband_records(self) -> pl.DataFrame:
        records = pl.read_csv(self.cytoband_file, separator="\t", has_header=False, new_columns=["chromosome", "start", "end", "name", "stain"])
        return (
            records.with_columns(pl.col("stain").replace(Stain2Color).alias("color"))
            .with_columns(pl.col("chromosome").map_elements(lambda x: is_active_chrom(x, self.gender), pl.Boolean).alias("active"))
            .filter(pl.col("active").eq(True))
            .drop("active")
        )

    @cached_property
    def chrom_records(self) -> pl.DataFrame:
        records = self.cytoband_records.group_by(["chromosome"], maintain_order=True).agg(pl.col("start").min(), pl.col("end").max())
        records = records.with_columns(
            pl.col("chromosome").str.replace("[Cc][Hh][Rr]", "").replace(["X", "Y"], ["23", "24"]).cast(pl.Int64).alias("c_order")
        ).sort("c_order")
        end_series = records["end"].cum_sum()
        start_series = pl.Series([0] + end_series[:-1].to_list()).rename("start")
        color_series = pl.Series((["#808080", "#E6E6E6"] * records.shape[0])[: records.shape[0]]).rename("color")
        return records.with_columns(start_series, end_series, color_series)

    def plot_genome_band(self, y: float):
        ticks = ((self.chrom_records["start"] + self.chrom_records["end"]) / 2).cast(pl.Int64)
        labels = self.chrom_records["chromosome"].str.replace("[Cc][Hh][Rr]", "")
        plt.xticks(ticks=ticks, labels=labels, size=7.5)
        for row in self.chrom_records.iter_rows(named=True):
            plt.hlines(y=y, xmin=row["start"], xmax=row["end"], lw=12, color=row["color"])

    def plot_chromosome_band(self, chrom: str, y: float):
        cytobands = self.cytoband_records.filter(pl.col("chromosome").eq(chrom))
        plt.xticks(ticks=(cytobands["start"] + cytobands["end"]) / 2, labels=cytobands["name"], rotation=90, size=6)
        for arm in ["p", "q"]:
            arm_cytobands = cytobands.filter(pl.col("name").str.starts_with(arm))
            lw = 12
            extra = arm_cytobands["end"].max() * 0.0068
            centromeric = [arm_cytobands["start"][0] + extra, arm_cytobands["end"][1], arm_cytobands["color"][0]]
            plt.hlines(y=y, xmin=centromeric[0], xmax=centromeric[1], lw=lw, capstyle="round", color=centromeric[2])
            centromeric = [arm_cytobands["start"][-2], arm_cytobands["end"][-1] - extra, arm_cytobands["color"][-1]]
            plt.hlines(y=y, xmin=centromeric[0], xmax=centromeric[1], lw=lw, capstyle="round", color=centromeric[2])
            for row in arm_cytobands[1:-1].iter_rows(named=True):
                plt.hlines(y=y, xmin=row["start"], xmax=row["end"], lw=lw, color=row["color"])

    def plot_gene_band(self, gene_records: pl.DataFrame, y: float):
        plt.xticks(ticks=(gene_records["tx_start"] + gene_records["tx_end"]) / 2, labels=gene_records["gene"], size=8)
        for row in gene_records.iter_rows(named=True):
            plt.plot((row["tx_start"], row["tx_end"]), (y, y), color=Color.GRAY.value)
            if row["strand"] == "+":
                plt.annotate("", xy=(row["tx_end"], y), xytext=(row["tx_start"], y), arrowprops=dict(arrowstyle="->", color=Color.GRAY.value))
            else:
                plt.annotate("", xy=(row["tx_start"], y), xytext=(row["tx_end"], y), arrowprops=dict(arrowstyle="<-", color=Color.GRAY.value))
            for start, end in zip(row["start"], row["end"]):
                plt.hlines(y=y, xmin=start, xmax=end, lw=12, color=Color.GRAY.value)

    def plot_genome_copynumber(self, zip_writer: zipfile.ZipFile):
        logger.info("Draw Genome")
        filtered_cn_records = self.mask_cn_records.filter(pl.col("cn").le(6))
        title = f"{self.sample} genome"
        ylim_max = math.ceil(max(4, filtered_cn_records["cn"].max()))
        plt.figure(figsize=self.figsize)
        plt.title(title, loc="center")
        plt.xlabel("Chromosome")
        plt.ylabel("Copy Number")
        plt.yticks(ticks=list(range(0, ylim_max)))
        plt.ylim(-2, ylim_max)
        for index, row in enumerate(self.chrom_records.iter_rows(named=True)):
            chrom = row["chromosome"]
            cn_records = filtered_cn_records.filter(pl.col("chromosome").eq(chrom))
            color = str([Color.BLUE.value, Color.ORANGE.value][index % 2])
            plt.scatter(((cn_records["start"] + cn_records["end"]) / 2) + row["start"], cn_records["cn"], s=3, c=color)
            plt.hlines(y=1 if is_male(self.gender) and is_sex_chrom(chrom) else 2, xmin=row["start"], xmax=row["end"], ls=":", color=Color.GRAY.value)
        self.plot_genome_band(y=-1.1)
        bytes_io = io.BytesIO()
        plt.savefig(bytes_io, dpi=self.dpi, bbox_inches="tight")
        plt.close()
        zip_writer.writestr(f"{self.sample}.genome.cnv.png", bytes_io.getvalue())

    def plot_chromosomes_copynumber(self, zip_writer: zipfile.ZipFile):
        filtered_cn_records = self.mask_cn_records.filter(pl.col("cn").le(6))
        for chrom in self.chrom_records["chromosome"]:
            logger.info("Draw %s", chrom)
            title = f"{self.sample} {chrom}"
            cn_records = filtered_cn_records.filter(pl.col("chromosome").eq(chrom))
            cytoband_records = self.cytoband_records.filter(pl.col("chromosome").eq(chrom))
            ylim_max = math.ceil(max(4, cn_records["cn"].max()))
            plt.figure(figsize=self.figsize)
            plt.title(title, loc="center")
            plt.xlabel("Location")
            plt.ylabel("Copy Number")
            plt.yticks(ticks=list(range(0, ylim_max)))
            plt.ylim(-2, ylim_max)
            plt.scatter((cn_records["start"] + cn_records["end"]) / 2, cn_records["cn"], s=3, c=Color.ORANGE.value)
            plt.hlines(
                y=1 if is_male(self.gender) and is_sex_chrom(chrom) else 2,
                xmin=cytoband_records["start"].min(),
                xmax=cytoband_records["end"].max(),
                ls=":",
                color=Color.GRAY.value,
            )
            self.plot_chromosome_band(chrom=chrom, y=-1.1)
            bytes_io = io.BytesIO()
            plt.savefig(bytes_io, dpi=self.dpi, bbox_inches="tight")
            plt.close()
            zip_writer.writestr(f"{self.sample}.{chrom}.cnv.png", bytes_io.getvalue())

    def plot_cnv_copynumber(self, zip_writer: zipfile.ZipFile):
        with VariantFile(str(self.call_vcf_file)) as reader:
            for cnv in reader:
                gene_records = self.gene_records.filter(
                    pl.col("chromosome").eq(cnv.chrom) & pl.col("tx_start").le(cnv.stop) & pl.col("tx_end").ge(cnv.start)
                )
                if not gene_records.is_empty():
                    key = f"{cnv.chrom}:{cnv.start}-{cnv.stop}"
                    logger.info("Plot %s", key)
                    title = f"{self.sample} {key}"
                    img = f"{self.sample}.{re.sub(r':|-', '_', key)}.cnv.png"
                    tx_start, tx_end = min(gene_records["tx_start"].min(), cnv.start), max(gene_records["tx_end"].max(), cnv.stop)
                    gene_cn_records = self.multi_cn_records.filter(
                        pl.col("chromosome").eq(cnv.chrom) & pl.col("start").le(tx_end) & pl.col("end").ge(tx_start)
                    )
                    cnv_cn_records = self.multi_cn_records.filter(
                        pl.col("chromosome").eq(cnv.chrom) & pl.col("start").le(cnv.stop) & pl.col("end").ge(cnv.start)
                    )
                    ylim_max = math.ceil(min(max(4, gene_cn_records[:, 3:].max_horizontal().max()), 7))
                    plt.figure(figsize=self.figsize)
                    plt.title(title, loc="center")
                    plt.xlabel("Location")
                    plt.ylabel("Position")
                    plt.yticks(ticks=list(range(0, ylim_max)))
                    plt.ylim(-2, ylim_max)
                    self.plot_gene_band(gene_records=gene_records, y=-1.1)
                    samples = self.multi_cn_records.columns[3:]
                    if len(samples) == 1:
                        plt.hlines(
                            y=1 if is_male(self.gender) and is_sex_chrom(cnv.chrom) else 2,
                            xmin=gene_cn_records["start"].min(),
                            xmax=gene_cn_records["end"].max(),
                            ls="dotted",
                            color=Color.GRAY.value,
                        )
                    for sample in samples[::-1]:
                        if sample == self.sample:
                            plt.plot(
                                (gene_cn_records["start"] + gene_cn_records["end"]) / 2,
                                gene_cn_records[sample].clip(upper_bound=7),
                                marker=".",
                                ms=9,
                                color=Color.GRAY.value,
                            )
                            plt.plot(
                                (cnv_cn_records["start"] + cnv_cn_records["end"]) / 2,
                                cnv_cn_records[sample].clip(upper_bound=7),
                                marker=".",
                                ms=9,
                                color=(Color.RED.value if cnv.alts[0] == "<DEL>" else Color.BLUE.value),
                            )
                        else:
                            plt.plot(
                                (gene_cn_records["start"] + gene_cn_records["end"]) / 2,
                                gene_cn_records[sample].clip(upper_bound=7),
                                marker=".",
                                ms=6,
                                color=Color.GRAY2.value,
                            )
                    bytes_io = io.BytesIO()
                    plt.savefig(bytes_io, dpi=self.dpi, bbox_inches="tight")
                    plt.close()
                    zip_writer.writestr(img, bytes_io.getvalue())

    def do_image(self):
        with zipfile.ZipFile(self.image_file, "w", zipfile.ZIP_DEFLATED) as zip_writer:
            self.plot_genome_copynumber(zip_writer=zip_writer)
            self.plot_chromosomes_copynumber(zip_writer=zip_writer)
            self.plot_cnv_copynumber(zip_writer=zip_writer)

    def run(self):
        if not self.image_file.exists():
            self.do_image()
